package io.daocloud.prometheustestdemo.interceptor;

import io.daocloud.prometheustestdemo.service.RateLimitService;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * 限流拦截器
 */
@Component
@Slf4j
public class RateLimitInterceptor implements HandlerInterceptor {
    
    //private static final Logger log = LoggerFactory.getLogger(RateLimitInterceptor.class);
    
    @Autowired
    private RateLimitService rateLimitService;

    private static final String GLOBAL_RATE_LIMIT_KEY = "global_rate_limit";
    
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        // 使用全局限流键实现统一限流
        String rateLimitKey = GLOBAL_RATE_LIMIT_KEY;
        // 检查是否允许请求
        if (!rateLimitService.isAllowed(rateLimitKey)) {
            log.warn("Rate limit exceeded for request: {} from IP: {}", request.getRequestURI(), getClientIP(request));
            response.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            response.setContentType("application/json;charset=UTF-8");
            response.getWriter().write("{\"error\":\"Too Many Requests\",\"message\":\"Rate limit exceeded\"}");
            return false;
        }
        
        return true;
    }

    /**
     * 获取客户端真实IP地址
     */
    private String getClientIP(HttpServletRequest request) {
        String xForwardedFor = request.getHeader("X-Forwarded-For");
        if (xForwardedFor != null && !xForwardedFor.isEmpty()) {
            return xForwardedFor.split(",")[0].trim();
        }
        
        String xRealIP = request.getHeader("X-Real-IP");
        if (xRealIP != null && !xRealIP.isEmpty()) {
            return xRealIP;
        }
        
        return request.getRemoteAddr();
    }
} 